package com.liuhengjia.springaiapply.vector;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;

import java.util.List;

/**
 * 文件工具类
 */
public class Files {
    /**
     * 加载 PDF 至向量数据库
     *
     * @return List<Document>
     */
    public static List<Document> getDocumentsFromPdf() {
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(
//                Files.class.getClassLoader().getResource("Mybatis.pdf").getPath(),
                "classpath:/Mybatis.pdf",
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)// 设置页面顶部边距为 0
                        .withPagesPerDocument(1)// 每个文档只包含一页内容
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)// 设置文本格式化器，删除顶部若干行（这里设置为 0 行）
                                .build())
                        .build()
        );
        return pagePdfDocumentReader.read();
    }

    /**
     * 加载 MD 文件至向量数据库
     *
     * @return List<Document>
     */
    public static List<Document> getDocumentsFormMarkdown() {
        MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(
//                Files.class.getClassLoader().getResource("Java基础.md").getPath()
                "classpath:/Java基础.md"
        );
        return markdownDocumentReader.read();
    }
}
